package com.agentframework.core.tools;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.lang.NonNull;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Stream;

/**
 * @author wangjun
 **/
public final class ToolManager implements ApplicationContextAware, InitializingBean {

    private static final Logger logger = LoggerFactory.getLogger(ToolManager.class);

    static List<ToolCallback> toolCallbackList = Collections.synchronizedList(new ArrayList<>());

    static Map<String, ToolCallback> toolCallbackMap = new ConcurrentHashMap<>();

    private ApplicationContext applicationContext;

    @Override
    public void afterPropertiesSet() throws Exception {
        Map<String, ToolCallbackProvider> toolCallbackProviderMapOfBeanMap = applicationContext.getBeansOfType(ToolCallbackProvider.class);
        ToolCallbackProvider[] toolCallbackProviders = toolCallbackProviderMapOfBeanMap.values().toArray(new ToolCallbackProvider[] {});
        List<ToolCallback[]> collect = Stream.of(toolCallbackProviders).map(ToolCallbackProvider::getToolCallbacks).toList();
        Map<String, AgentTool> agentToolMap = applicationContext.getBeansOfType(AgentTool.class);
        ToolCallback[] toolCallbacks = ToolCallbacks.from(agentToolMap.values().toArray());
        logger.debug("tools: {}", toolCallbacks.length);
        collect.forEach(item -> toolCallbackList.addAll(List.of(item)));
        toolCallbackList.addAll(List.of(toolCallbacks));
        toolCallbackList.forEach(toolCallback -> toolCallbackMap.put(toolCallback.getToolDefinition().name(), toolCallback));
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    public static List<ToolCallback> getTools() {
        return toolCallbackList;
    }

    public static ToolCallback getTool(String toolName) {
        return toolCallbackMap.get(toolName);
    }

}
